package code.gateplugins.libsvm;

import gate.creole.ExecutionException;
import gate.creole.ResourceInstantiationException;

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.FileWriter;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.net.URISyntaxException;
import java.net.URL;
import java.util.HashMap;

import org.apache.commons.io.IOUtils;

import libsvm.svm;
import libsvm.svm_node;
import libsvm.svm_problem;

public class ModelIO {
	
	static final String SEP = System.getProperty("file.separator");
	static final String ALPHABET_SUFFIX = SEP + "libsvm.alphabet";
	static final String MODEL_SUFFIX = SEP + "libsvm.model";
	static final String DUMP_SUFFIX = SEP + "libsvm.dump";

	public static SvmData load(URL path, boolean alphabetsOnly) throws ResourceInstantiationException {
		ObjectInputStream ais = null;
		SvmData svmData = new SvmData();
		try {
			String f = getTemplate(path);
			if(!alphabetsOnly) {
				svmData.setModel(svm.svm_load_model(f + MODEL_SUFFIX));
			}
			
			ais = new ObjectInputStream(new FileInputStream(f + ALPHABET_SUFFIX));
			@SuppressWarnings("unchecked")
			HashMap<String, Integer> lAlphabet = (HashMap<String, Integer>) ais.readObject();
			svmData.setLabelAlphabet(lAlphabet);
			@SuppressWarnings("unchecked")
			HashMap<String, Integer> alphabet = (HashMap<String, Integer>) ais.readObject();
			svmData.setAlphabet(alphabet);
		} catch (IOException | ClassNotFoundException | URISyntaxException e) {
			throw new ResourceInstantiationException(e);
		} finally {
			IOUtils.closeQuietly(ais);
		}
		return svmData;
	}
	
	public static void dump(URL modelFile, svm_problem problem, 
			SvmData svmData, boolean continuousEval) throws ExecutionException {
		BufferedWriter writer = null;
		String f = null;
		try {
			f = getTemplate(modelFile);
			String dumpFile = f + DUMP_SUFFIX;
			writer = new BufferedWriter(new FileWriter(dumpFile, continuousEval)); // if continuous dump - append is true
			for (int i = 0; i < problem.l; ++i) {
				writer.write(problem.y[i] + " ");
				svm_node[] x = problem.x[i];
				for (int j = 0; j < x.length - 1; ++j) {
					writer.write(x[j].index + ":" + x[j].value + " ");
				}
				writer.write(x[x.length - 1].index + ":" + x[x.length - 1].value + "\n");
			}
		} catch (IOException | URISyntaxException e) {
			throw new ExecutionException(e);
		} finally {
			IOUtils.closeQuietly(writer);
		}
		// if continuous alphabets are alredy fixed
		if(!continuousEval) {
			storeAlphabets(svmData, f);
		}
	}

	public static void save(URL modelPath, SvmData svmData) throws ExecutionException {
		String f = null;
		try {
			f = getTemplate(modelPath);
			svm.svm_save_model(f + MODEL_SUFFIX, svmData.getModel());
		} catch (URISyntaxException | IOException e) {
			throw new ExecutionException(e);
		}
		
		storeAlphabets(svmData, f);
	}

	private static void storeAlphabets(SvmData svmData, String f) throws ExecutionException {
		ObjectOutputStream aos = null;
		try {
			aos = new ObjectOutputStream(new FileOutputStream(f + ALPHABET_SUFFIX));
			aos.writeObject(svmData.getLabelAlphabet());
			aos.writeObject(svmData.getAlphabet());
		} catch (IOException e) {
			throw new ExecutionException(e);
		} finally {
			IOUtils.closeQuietly(aos);
		}
	}

	private static String getTemplate(URL modelPath) throws URISyntaxException {
		File f = new File(modelPath.toURI());
		f.mkdirs();
		return f.getAbsolutePath();
	}
}
